// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>
#include <memory>
#include <string>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>

#include <nebula/array/data.h>
#include <nebula/types/scalar.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/checked_cast.h>
#include <turbo/base/macros.h>

namespace nebula {

    class Array;

    class ChunkedArray;

    class RecordBatch;

    class Table;

    /// \class Datum
    /// \brief Variant type for various Nebula C++ data structures
    struct TURBO_EXPORT Datum {
        /// \brief The kind of datum stored
        enum Kind {
            NONE, SCALAR, ARRAY, CHUNKED_ARRAY, RECORD_BATCH, TABLE
        };

        /// \brief A placeholder type to represent empty datum
        struct Empty {
        };

        /// \brief Datums variants may have a length. This special value indicate that the
        /// current variant does not have a length.
        static constexpr int64_t kUnknownLength = -1;

        /// \brief Storage of the actual datum.
        ///
        /// Note: For arrays, ArrayData is stored instead of Array for easier processing
        std::variant<Empty, std::shared_ptr<Scalar>, std::shared_ptr<ArrayData>,
                std::shared_ptr<ChunkedArray>, std::shared_ptr<RecordBatch>,
                std::shared_ptr<Table>>
                value;

        /// \brief Empty datum, to be populated elsewhere
        Datum() = default;

        Datum(const Datum &other) = default;

        Datum &operator=(const Datum &other) = default;

        Datum(Datum &&other) = default;

        Datum &operator=(Datum &&other) = default;

        /// \brief Construct from a Scalar
        Datum(std::shared_ptr<Scalar> value)  // NOLINT implicit conversion
                : value(std::move(value)) {}

        /// \brief Construct from an ArrayData
        Datum(std::shared_ptr<ArrayData> value)  // NOLINT implicit conversion
                : value(std::move(value)) {}

        /// \brief Construct from an ArrayData
        Datum(ArrayData arg)  // NOLINT implicit conversion
                : value(std::make_shared<ArrayData>(std::move(arg))) {}

        /// \brief Construct from an Array
        Datum(const Array &value);  // NOLINT implicit conversion

        /// \brief Construct from an Array
        Datum(const std::shared_ptr<Array> &value);  // NOLINT implicit conversion

        /// \brief Construct from a ChunkedArray
        Datum(std::shared_ptr<ChunkedArray> value);  // NOLINT implicit conversion

        /// \brief Construct from a RecordBatch
        Datum(std::shared_ptr<RecordBatch> value);  // NOLINT implicit conversion

        /// \brief Construct from a Table
        Datum(std::shared_ptr<Table> value);  // NOLINT implicit conversion

        /// \brief Construct from a ChunkedArray.
        ///
        /// This can be expensive, prefer the shared_ptr<ChunkedArray> constructor
        explicit Datum(const ChunkedArray &value);

        /// \brief Construct from a RecordBatch.
        ///
        /// This can be expensive, prefer the shared_ptr<RecordBatch> constructor
        explicit Datum(const RecordBatch &value);

        /// \brief Construct from a Table.
        ///
        /// This can be expensive, prefer the shared_ptr<Table> constructor
        explicit Datum(const Table &value);

        /// \brief Cast from concrete subtypes of Array or Scalar to Datum
        template<typename T, bool IsArray = std::is_base_of_v<Array, T>,
                bool IsScalar = std::is_base_of_v<Scalar, T>,
                typename = enable_if_t<IsArray || IsScalar>>
        Datum(std::shared_ptr<T> value)  // NOLINT implicit conversion
                : Datum(std::shared_ptr<typename std::conditional<IsArray, Array, Scalar>::type>(
                std::move(value))) {}

        /// \brief Cast from concrete subtypes of Array or Scalar to Datum
        template<typename T, typename TV = typename std::remove_reference_t<T>,
                bool IsArray = std::is_base_of_v<Array, T>,
                bool IsScalar = std::is_base_of_v<Scalar, T>,
                typename = enable_if_t<IsArray || IsScalar>>
        Datum(T &&value)  // NOLINT implicit conversion
                : Datum(std::make_shared<TV>(std::forward<T>(value))) {}

        /// \brief Copy from concrete subtypes of Scalar.
        ///
        /// The concrete scalar type must be copyable (not all of them are).
        template<typename T, typename = enable_if_t<std::is_base_of_v<Scalar, T>>>
        Datum(const T &value)  // NOLINT implicit conversion
                : Datum(std::make_shared<T>(value)) {}

        // Convenience constructors
        /// \brief Convenience constructor storing a bool scalar.
        explicit Datum(bool value);

        /// \brief Convenience constructor storing an int8 scalar.
        explicit Datum(int8_t value);

        /// \brief Convenience constructor storing a uint8 scalar.
        explicit Datum(uint8_t value);

        /// \brief Convenience constructor storing an int16 scalar.
        explicit Datum(int16_t value);

        /// \brief Convenience constructor storing a uint16 scalar.
        explicit Datum(uint16_t value);

        /// \brief Convenience constructor storing an int32 scalar.
        explicit Datum(int32_t value);

        /// \brief Convenience constructor storing a uint32 scalar.
        explicit Datum(uint32_t value);

        /// \brief Convenience constructor storing an int64 scalar.
        explicit Datum(int64_t value);

        /// \brief Convenience constructor storing a uint64 scalar.
        explicit Datum(uint64_t value);

        /// \brief Convenience constructor storing a float scalar.
        explicit Datum(float value);

        /// \brief Convenience constructor storing a double scalar.
        explicit Datum(double value);

        /// \brief Convenience constructor storing a string scalar.
        explicit Datum(std::string value);

        /// \brief Convenience constructor storing a string scalar.
        explicit Datum(const char *value);

        /// \brief Convenience constructor for a DurationScalar from std::chrono::duration
        template<template<typename, typename> class StdDuration, typename Rep,
                typename Period,
                typename = decltype(DurationScalar{StdDuration<Rep, Period>{}})>
        explicit Datum(StdDuration<Rep, Period> d) : Datum{DurationScalar(d)} {}

        /// \brief The kind of data stored in Datum
        Datum::Kind kind() const {
            switch (this->value.index()) {
                case 0:
                    return Datum::NONE;
                case 1:
                    return Datum::SCALAR;
                case 2:
                    return Datum::ARRAY;
                case 3:
                    return Datum::CHUNKED_ARRAY;
                case 4:
                    return Datum::RECORD_BATCH;
                case 5:
                    return Datum::TABLE;
                default:
                    return Datum::NONE;
            }
        }

        /// \brief Retrieve the stored array as ArrayData
        ///
        /// Use make_array() if an Array is desired (which is more expensive).
        /// \throws std::bad_variant_access if the datum is not an array
        const std::shared_ptr<ArrayData> &array() const {
            return std::get<std::shared_ptr<ArrayData>>(this->value);
        }

        /// \brief The sum of bytes in each buffer referenced by the datum
        /// Note: Scalars report a size of 0
        /// \see nebula::util::TotalBufferSize for caveats
        int64_t TotalBufferSize() const;

        /// \brief Get the stored ArrayData in mutable form
        ///
        /// For internal use primarily. Keep in mind a shared_ptr<Datum> may have multiple
        /// owners.
        ArrayData *mutable_array() const { return this->array().get(); }

        /// \brief Retrieve the stored array as Array
        /// \throws std::bad_variant_access if the datum is not an array
        std::shared_ptr<Array> make_array() const;

        /// \brief Retrieve the chunked array stored
        /// \throws std::bad_variant_access if the datum is not a chunked array
        const std::shared_ptr<ChunkedArray> &chunked_array() const {
            return std::get<std::shared_ptr<ChunkedArray>>(this->value);
        }

        /// \brief Retrieve the record batch stored
        /// \throws std::bad_variant_access if the datum is not a record batch
        const std::shared_ptr<RecordBatch> &record_batch() const {
            return std::get<std::shared_ptr<RecordBatch>>(this->value);
        }

        /// \brief Retrieve the table stored
        /// \throws std::bad_variant_access if the datum is not a table
        const std::shared_ptr<Table> &table() const {
            return std::get<std::shared_ptr<Table>>(this->value);
        }

        /// \brief Retrieve the scalar stored
        /// \throws std::bad_variant_access if the datum is not a scalar
        const std::shared_ptr<Scalar> &scalar() const {
            return std::get<std::shared_ptr<Scalar>>(this->value);
        }

        /// \brief Retrieve the datum as its concrete array type
        /// \throws std::bad_variant_access if the datum is not an array
        /// \tparam ExactType the expected array type, may cause undefined behavior if it is not
        /// the type of the stored array
        template<typename ExactType>
        std::shared_ptr<ExactType> array_as() const {
            return turbo::checked_pointer_cast<ExactType>(this->make_array());
        }

        /// \brief Retrieve the datum as its concrete scalar type
        /// \throws std::bad_variant_access if the datum is not a scalar
        /// \tparam ExactType the expected scalar type, may cause undefined behavior if it is
        /// not the type of the stored scalar
        template<typename ExactType>
        const ExactType &scalar_as() const {
            return turbo::checked_cast<const ExactType &>(*this->scalar());
        }

        /// \brief True if Datum contains an array
        bool is_array() const { return this->kind() == Datum::ARRAY; }

        /// \brief True if Datum contains a chunked array
        bool is_chunked_array() const { return this->kind() == Datum::CHUNKED_ARRAY; }

        /// \brief True if Datum contains an array or a chunked array
        bool is_arraylike() const {
            return this->kind() == Datum::ARRAY || this->kind() == Datum::CHUNKED_ARRAY;
        }

        /// \brief True if Datum contains a scalar
        bool is_scalar() const { return this->kind() == Datum::SCALAR; }

        /// \brief True if Datum contains a scalar or array-like data
        bool is_value() const { return this->is_arraylike() || this->is_scalar(); }

        /// \brief Return the null count.
        ///
        /// Only valid for scalar and array-like data.
        int64_t null_count() const;

        /// \brief The value type of the variant, if any
        ///
        /// \return nullptr if no type
        const std::shared_ptr<DataType> &type() const;

        /// \brief The schema of the variant, if any
        ///
        /// \return nullptr if no schema
        const std::shared_ptr<Schema> &schema() const;

        /// \brief The value length of the variant, if any
        ///
        /// \return kUnknownLength if no type
        int64_t length() const;

        /// \brief The array chunks of the variant, if any
        ///
        /// \return empty if not arraylike
        ArrayVector chunks() const;

        /// \brief True if the two data are equal
        bool equals(const Datum &other) const;

        bool operator==(const Datum &other) const { return equals(other); }

        bool operator!=(const Datum &other) const { return !equals(other); }

        std::string to_string() const;
    };

    TURBO_EXPORT void print_to(const Datum &, std::ostream *);

    TURBO_EXPORT std::string to_string(Datum::Kind kind);

}  // namespace nebula
